from data_type import *
from datetime import *

NP_Table = {'ftp': '21', 'ssh': '22', 'telnet': '23', 'smtp': '25', 'dns': '53', 'domain': '53', 'bootps': '67', 'bootpc': '68',
            'http': '80', 'pop3': '110', 'netbios_ns': '137', 'netbios_dgm': '138',
            'netbios_ssn': '139', 'https': '443', 'ssdp': '1900', 'msnp': '666', 'llmnr': '5355', 'mdns': '5353'} #msn是什么协议？

def GetValue(protInfoList:list, index:int, key:str):  #输入协议信息列表（格式见上）、查询列表下标和查询对象的键，返回查询对象的值
    objStr = protInfoList[index]
    if key not in objStr:
        return None
    startIndex = objStr.find(key) + len(key) + 1     #到'='之后
    if key == 'load':
        endIndex = objStr.find(key) + objStr[objStr.find(key):].find('"')
    else:
        endIndex = objStr.find(key) + objStr[objStr.find(key):].find(' ') #键开始的下标+键值对长度=值结束的下标
    return objStr[startIndex:endIndex]

def GetEthernetInfo(protInfoList:list):
    time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    dst = GetValue(protInfoList, 0, 'dst')
    src = GetValue(protInfoList, 0, 'src')
    type_ = GetValue(protInfoList, 0, 'type')
    return time, dst, src, type_
 
def GetIpv4Info(protInfoList:list):
    time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    version = int(GetValue(protInfoList, 1, 'version'))
    headerLen = int(GetValue(protInfoList, 1, 'ihl'))
    tos = int(GetValue(protInfoList, 1, 'tos'), base=16)
    packetLen = int(GetValue(protInfoList, 1, 'len'))
    protoId = int(GetValue(protInfoList, 1, 'id'))
    flags = GetValue(protInfoList, 1, 'flags')
    frag = int(GetValue(protInfoList, 1, 'frag'))
    ttl = int(GetValue(protInfoList, 1, 'ttl'))
    proto = GetValue(protInfoList, 1, 'proto')
    chksum = int(GetValue(protInfoList, 1, 'chksum'), base=16)
    srcIp = GetValue(protInfoList, 1, 'src')
    dstIp = GetValue(protInfoList, 1, 'dst')
    return time, version, headerLen, tos, packetLen, protoId, flags, frag, ttl, proto, chksum, srcIp, dstIp

def GetIpv6Info(protInfoList:list):
    time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    version = int(GetValue(protInfoList, 1, 'version'))
    tc = int(GetValue(protInfoList, 1, 'tc'))
    fl = int(GetValue(protInfoList, 1, 'fl'))
    plen = int(GetValue(protInfoList, 1, 'plen'))
    nh = GetValue(protInfoList, 1, 'nh')
    hlim = int(GetValue(protInfoList, 1, 'hlim'))
    srcIp = GetValue(protInfoList, 1, 'src')
    dstIp = GetValue(protInfoList, 1, 'dst')
    return time, version, tc, fl, plen, nh, hlim, srcIp, dstIp
    
def GetTcpInfo(protInfoList:list):
    time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    sport = GetValue(protInfoList, 2, 'sport')
    if sport in NP_Table:
        sport = NP_Table.get(sport)
    dport = GetValue(protInfoList, 2, 'dport')
    if dport in NP_Table:
        dport = NP_Table.get(dport)
    seq = int(GetValue(protInfoList, 2, 'seq'))
    ack = int(GetValue(protInfoList, 2, 'ack'))
    dataofs = int(GetValue(protInfoList, 2, 'dataofs'))
    reserved = int(GetValue(protInfoList, 2, 'reserved'))
    flags = GetValue(protInfoList, 2, 'flags')
    window = int(GetValue(protInfoList, 2, 'window'))
    chksum = int(GetValue(protInfoList, 2, 'chksum'), base=16)
    urgptr = int(GetValue(protInfoList, 2, 'urgptr'))
    startIndex = protInfoList[3].find('load') + len('load') + 1
    raw = protInfoList[3][startIndex:]
    return time, sport, dport, seq, ack, dataofs, reserved, flags, window, chksum, urgptr, raw

def GetUdpInfo(protInfoList:list):
    time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    sport = GetValue(protInfoList, 2, 'sport')
    if sport in NP_Table:
        sport = NP_Table.get(sport)
    dport = GetValue(protInfoList, 2, 'dport')
    if dport in NP_Table:
        dport = NP_Table.get(dport)
    packetLen = int(GetValue(protInfoList, 2, 'len'))
    chksum = int(GetValue(protInfoList, 2, 'chksum'), base=16)
    return time, sport, dport, packetLen, chksum

def GetArpInfo(protInfoList:list):
    time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    hwtype = GetValue(protInfoList, 1, 'hwtype')
    ptype = GetValue(protInfoList, 1, 'ptype')
    hwlen = int(GetValue(protInfoList, 1, 'hwlen'))
    plen = int(GetValue(protInfoList, 1, 'plen'))
    op = GetValue(protInfoList, 1, 'op')
    hwsrc = GetValue(protInfoList, 1, 'hwsrc')
    psrc = GetValue(protInfoList, 1, 'psrc')
    hwdst = GetValue(protInfoList, 1, 'hwdst')
    pdst = GetValue(protInfoList, 1, 'pdst')
    return time, hwtype, ptype, hwlen, plen, op, hwsrc, psrc, hwdst, pdst

def GetDnsInfo(protInfoList:list):
    time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    id = int(GetValue(protInfoList, 3, 'id'))
    flags = {}
    flagsKeys = ['qr','opcode','aa','tc','rd','ra','z','ad','cd','rcode']
    for i in flagsKeys:
        flags[i] = GetValue(protInfoList, 3, i)
        
    qdcount = int(GetValue(protInfoList, 3, 'qdcount'))
    ancount = int(GetValue(protInfoList, 3, 'ancount'))
    nscount = int(GetValue(protInfoList, 3, 'nscount'))
    arcount = int(GetValue(protInfoList, 3, 'arcount'))
    
    qdarea = {}
    qdKeys = ['qname', 'qtype', 'qclass']
    for i in qdKeys:
        qdarea[i] = GetValue(protInfoList, 3, i)
        
    anarea = []
    anKeys = ['rrname', 'type', 'rclass', 'ttl', 'rdlen', 'rdata'] #因为仅看了部分示例，字段可能不全，如果需要后续再补充
    for i in range(ancount):
        an = {}
        for j in anKeys:
            an[j] = GetValue(protInfoList, 4 + i, j)
        anarea.append(an)
        
    nsarea = []
    nsKeys = ['rrname', 'type', 'rclass', 'ttl', 'mname', 'rname', 'serial', 'refresh', 'retry', 'expire', 'minimum'] #因为仅看了部分示例，字段可能不全，如果需要后续再补充
    for i in range(nscount):
        ns = {}
        for j in nsKeys:
            ns[j] = GetValue(protInfoList, 4 + ancount + i, j)
        nsarea.append(ns)
    
    ararea = []
    arKeys = ['rrname', 'type', 'rclass', 'ectrcode', 'version', 'z', 'rdlen'] #因为仅看了部分示例，字段可能不全，如果需要后续再补充
    for i in range(arcount):
        ar = {}
        for j in arKeys:
            ar[j] = GetValue(protInfoList, 4 + ancount + nscount + i, j)
        ararea.append(ar)
    
    return time, id, flags, qdcount, ancount, nscount, arcount, qdarea, anarea, nsarea, ararea

def GetIcmpv4Info(protInfoList:list):
    time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    type_ = GetValue(protInfoList, 2, 'type')
    code = GetValue(protInfoList, 2, 'code')
    chksum = int(GetValue(protInfoList, 2, 'chksum'), base=16)
    reserved = GetValue(protInfoList, 2, 'reserved')
    if reserved != None:
        reserved = int(reserved)
    length = GetValue(protInfoList, 2, 'length')
    if length != None:
        length = int(length)
    nexthopmtu = GetValue(protInfoList, 2, 'nexthopmtu')
    if nexthopmtu != None:
        nexthopmtu = int(nexthopmtu)
    unused = GetValue(protInfoList, 2, 'unused')
    IpError = {}
    IpErrorKeys = ['version','ihl','tos','len','id','flags','frag','ttl','proto','chksum','src','dst']
    for i in IpErrorKeys:
        IpError[i] = GetValue(protInfoList, 3, i)
        
    UdpError = {}
    UdpErrorKeys = ['sport','dport','len','chksum']
    for i in UdpErrorKeys:
        UdpError[i] = GetValue(protInfoList, 4, i)
    return time, type_, code, chksum, reserved, length, nexthopmtu, unused, IpError, UdpError

def GetIcmpv6Info(protInfoList:list):
    time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    Icmpv6ND_RA = {}
    Icmpv6ND_RAKeys = ['type','code','cksum','chlim','M','O','H','prf','P','res','routerlifetime','reachabletime','retranstimer']
    for i in Icmpv6ND_RAKeys:
        Icmpv6ND_RA[i] = GetValue(protInfoList, 2, i)
    
    Icmpv6NDOptSrcLLAddr = {}
    if len(protInfoList) >= 4:
        Icmpv6NDOptSrcLLAddrKeys = ['type','len','lladd']
        for i in Icmpv6NDOptSrcLLAddrKeys:
            Icmpv6NDOptSrcLLAddr[i] = GetValue(protInfoList, 3, i)
    
    Icmpv6NDOptMTU = {}
    if len(protInfoList) >= 5:
        Icmpv6NDOptMTUKeys = ['type','len','res','mtu']
        for i in Icmpv6NDOptMTUKeys:
            Icmpv6NDOptMTU[i] = GetValue(protInfoList, 4, i)
    
    Icmpv6NDOptPrefixInfo = {}
    if len(protInfoList) >= 6:
        Icmpv6NDOptPrefixInfoKeys = ['type','len','perfixlen','L','A','R','res1','validlifetime','preferredlifetime','res2','prefix']
        for i in Icmpv6NDOptPrefixInfoKeys:
            Icmpv6NDOptPrefixInfo[i] = GetValue(protInfoList, 5, i)
    
    return time, Icmpv6ND_RA, Icmpv6NDOptSrcLLAddr, Icmpv6NDOptMTU, Icmpv6NDOptPrefixInfo

def GetHttpReqInfo(protInfoList:list):
    time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    method_ = GetValue(protInfoList, 4, 'Method')
    path = GetValue(protInfoList, 4, 'Path')
    httpVersion = GetValue(protInfoList, 4, 'Http_Version')
    accept = GetValue(protInfoList, 4, 'Accept')
    cacheControl = GetValue(protInfoList, 4, 'Cache_Control')
    connection = GetValue(protInfoList, 4, 'Connection')
    contentLength = GetValue(protInfoList, 4, 'Content_Length')
    # contentLength = GetValue(protInfoList, 4, 'Content_Length')
    # if contentLength != None:
    #     contentLength = contentLength[1:-1]
    #     contentLength = int(contentLength) 
    contentType = GetValue(protInfoList, 4, 'Content_Type')
    host = GetValue(protInfoList, 4, 'Host')
    upgrade = GetValue(protInfoList, 4, 'Upgrade')
    userAgent = GetValue(protInfoList, 4, 'User_Agent')
    unknownHeaders = GetValue(protInfoList, 4, 'Unknown_Headers')
    return time, method_, path, httpVersion, accept, cacheControl, connection, contentLength, contentType, host, upgrade, userAgent, unknownHeaders

def GetHttpResInfo(protInfoList:list):
    time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    httpVersion = GetValue(protInfoList, 4, 'Http_Version')
    statusCode = GetValue(protInfoList, 4, 'Accept')
    # if statusCode != None:
    #     statusCode = statusCode[1:-1]
    #     statusCode = int(statusCode)
    reasonPhrase = GetValue(protInfoList, 4, 'Cache_Control')
    connection = GetValue(protInfoList, 4, 'Connection')
    contentLength = GetValue(protInfoList, 4, 'Content_Length')
    # if contentLength != None:
    #     contentLength = contentLength[1:-1]
    #     contentLength = int(contentLength)
    contentType = GetValue(protInfoList, 4, 'Content_Type')
    return time, httpVersion, statusCode, reasonPhrase, connection, contentLength, contentType

def GetOnePkt(protocolInfoList: list, pkt: Packet):
    #pkt.id = count
    if '<Ether' in protocolInfoList[0]:
        pkt.flag['ethernet'] = 1
        Ethernet = EthernetPkt()
        EthernetInfo = GetEthernetInfo(protocolInfoList)  #元组
        Ethernet.time = EthernetInfo[0]
        Ethernet.dst = EthernetInfo[1]
        Ethernet.src = EthernetInfo[2]
        Ethernet.type_ = EthernetInfo[3]
        #Ethernet.PrintMyInfo()
        pkt.ethernet = Ethernet   #这样写可不可以
        
    if '<TCP' in protocolInfoList[2]:
        pkt.flag['tcp'] = 1
        Tcp = TcpPkt()
        TcpInfo = GetTcpInfo(protocolInfoList)
        Tcp.time = TcpInfo[0]
        Tcp.sport = TcpInfo[1]
        Tcp.dport = TcpInfo[2]
        Tcp.seq = TcpInfo[3]
        Tcp.ack = TcpInfo[4]
        Tcp.dataofs = TcpInfo[5]
        Tcp.reserved = TcpInfo[6]
        Tcp.flags = TcpInfo[7]
        Tcp.window = TcpInfo[8]
        Tcp.chksum = TcpInfo[9]
        Tcp.urgptr = TcpInfo[10]
        Tcp.raw = TcpInfo[11]
        # Tcp.PrintMyInfo()
        pkt.tcp = Tcp
        
        if '<IPv6' in protocolInfoList[1]:
            pkt.flag['ipv6'] = 1
            Ipv6 = Ipv4Pkt()
            IpInfo = GetIpv6Info(protocolInfoList)
            Ipv6.time = IpInfo[0]
            Ipv6.version = IpInfo[1]
            Ipv6.tc = IpInfo[2] 
            Ipv6.fl = IpInfo[3] 
            Ipv6.plen = IpInfo[4] 
            Ipv6.nh = IpInfo[5] 
            Ipv6.hlim = IpInfo[6]  
            Ipv6.srcIp = IpInfo[7] 
            Ipv6.dstIp = IpInfo[8]
            #Ipv6.PrintMyInfo()
            pkt.ipv6 = Ipv6
            
        elif '<IP' in protocolInfoList[1]:
            pkt.flag['ipv4'] = 1
            Ipv4 = Ipv4Pkt()
            IpInfo = GetIpv4Info(protocolInfoList)
            Ipv4.time = IpInfo[0]
            Ipv4.version = IpInfo[1]
            Ipv4.headerLen = IpInfo[2] 
            Ipv4.tos = IpInfo[3] 
            Ipv4.packetLen = IpInfo[4] 
            Ipv4.protoId = IpInfo[5] 
            Ipv4.flags = IpInfo[6]  
            Ipv4.frag = IpInfo[7]  
            Ipv4.ttl = IpInfo[8] 
            Ipv4.proto = IpInfo[9] 
            Ipv4.chksum = IpInfo[10] 
            Ipv4.srcIp = IpInfo[11] 
            Ipv4.dstIp = IpInfo[12]
            #Ipv4.PrintMyInfo()
            pkt.ipv4 = Ipv4

        if len(protocolInfoList) >= 5:
            if '<HTTPRequest' in protocolInfoList[4]:
                pkt.flag['httpreq'] = 1
                HttpReq = HttpReqPkt()
                HTTPReqInfo = GetHttpReqInfo(protocolInfoList)
                HttpReq.time = HTTPReqInfo[0]
                HttpReq.method_ = HTTPReqInfo[1]
                HttpReq.path = HTTPReqInfo[2] 
                HttpReq.httpVersion = HTTPReqInfo[3] 
                HttpReq.accept = HTTPReqInfo[4] 
                HttpReq.cacheControl = HTTPReqInfo[5] 
                HttpReq.connection = HTTPReqInfo[6]  
                HttpReq.contentLength = HTTPReqInfo[7]  
                HttpReq.contentType = HTTPReqInfo[8] 
                HttpReq.host= HTTPReqInfo[9] 
                HttpReq.upgrade = HTTPReqInfo[10] 
                HttpReq.userAgent = HTTPReqInfo[11] 
                HttpReq.unknownHeaders = HTTPReqInfo[12]
                #HttpReq.PrintMyInfo()
                pkt.httpreq = HttpReq
                
            elif '<HTTPResponse' in protocolInfoList[4]:
                pkt.flag['httpres'] = 1
                HttpRes = HttpResPkt()
                HTTPResInfo = GetHttpResInfo(protocolInfoList)
                HttpRes.time = HTTPResInfo[0] 
                HttpRes.httpVersion = HTTPResInfo[1] 
                HttpRes.statusCode = HTTPResInfo[2] 
                HttpRes.reasonPhrase = HTTPResInfo[3] 
                HttpRes.connection = HTTPResInfo[4]  
                HttpRes.contentLength = HTTPResInfo[5]  
                HttpRes.contentType = HTTPResInfo[6] 
                #HttpRes.PrintMyInfo()
                pkt.httpres = HttpRes
        #if '<DNS' in protocolInfoList[3]:  TCP有吗？
        
    elif '<UDP' in protocolInfoList[2]:

        pkt.flag['udp'] = 1
        Udp = UdpPkt()
        UdpInfo = GetUdpInfo(protocolInfoList)
        Udp.time = UdpInfo[0]
        Udp.sport = UdpInfo[1]
        Udp.dport = UdpInfo[2]
        Udp.packetLen = UdpInfo[3]
        Udp.chksum = UdpInfo[4]
        #Udp.PrintMyInfo()
        pkt.udp = Udp
        
        if '<IPv6' in protocolInfoList[1]:  #会不会和tcp分支的变量冲突？
            pkt.flag['ipv6'] = 1
            Ipv6 = Ipv6Pkt()
            IpInfo = GetIpv6Info(protocolInfoList)
            Ipv6.time = IpInfo[0]
            Ipv6.version = IpInfo[1]
            Ipv6.tc = IpInfo[2] 
            Ipv6.fl = IpInfo[3] 
            Ipv6.plen = IpInfo[4] 
            Ipv6.nh = IpInfo[5] 
            Ipv6.hlim = IpInfo[6]  
            Ipv6.srcIp = IpInfo[7] 
            Ipv6.dstIp = IpInfo[8]
            #Ipv6.PrintMyInfo()
            pkt.ipv6 = Ipv6
            
        elif '<IP' in protocolInfoList[1]:
            pkt.flag['ipv4'] = 1
            Ipv4 = Ipv4Pkt()
            IpInfo = GetIpv4Info(protocolInfoList)
            Ipv4.time = IpInfo[0]
            Ipv4.version = IpInfo[1]
            Ipv4.headerLen = IpInfo[2] 
            Ipv4.tos = IpInfo[3] 
            Ipv4.packetLen = IpInfo[4] 
            Ipv4.protoId = IpInfo[5] 
            Ipv4.flags = IpInfo[6]  
            Ipv4.frag = IpInfo[7]  
            Ipv4.ttl = IpInfo[8] 
            Ipv4.proto = IpInfo[9] 
            Ipv4.chksum = IpInfo[10] 
            Ipv4.srcIp = IpInfo[11] 
            Ipv4.dstIp = IpInfo[12]
            #Ipv4.PrintMyInfo()
            pkt.ipv4 = Ipv4
            
        #if '<HTTPRequest' in protocolInfoList[4]:  UDP上有HTTP吗？  
        #elif '<HTTPResponse' in protocolInfoList[4]:
            
        if '<DNS' in protocolInfoList[3]:
            pkt.flag['dns'] = 1
            Dns = DnsPkt()
            DnsInfo = GetDnsInfo(protocolInfoList)
            Dns.time = DnsInfo[0]
            Dns.id = DnsInfo[1]
            Dns.flags = DnsInfo[2] 
            Dns.qdcount = DnsInfo[3] 
            Dns.ancount = DnsInfo[4] 
            Dns.nscount = DnsInfo[5] 
            Dns.arcount = DnsInfo[6]  
            Dns.qdarea = DnsInfo[7]  
            Dns.anarea = DnsInfo[8] 
            Dns.nsarea = DnsInfo[9] 
            Dns.ararea = DnsInfo[10]            
            #Dns.PrintMyInfo()
            pkt.dns = Dns
            
    elif '<ARP' in protocolInfoList[1]:
        pkt.flag['arp'] = 1
        Arp = ArpPkt()
        ArpInfo = GetArpInfo(protocolInfoList)
        Arp.time = ArpInfo[0]
        Arp.hwtype = ArpInfo[1]
        Arp.ptype = ArpInfo[2]
        Arp.hwlen = ArpInfo[3]
        Arp.plen = ArpInfo[4]
        Arp.op = ArpInfo[5]
        Arp.hwsrc = ArpInfo[6]
        Arp.psrc = ArpInfo[7]
        Arp.hwdst = ArpInfo[8]
        Arp.pdst = ArpInfo[9]
        #Arp.PrintMyInfo()
        pkt.arp = Arp
        
    elif '<ICMP' in protocolInfoList[2] or '<ICMPv6' in protocolInfoList[2]:
        if '<ICMPv6' in protocolInfoList[2]:
            pkt.flag['icmpv6'] = 1
            Icmpv6 = Icmpv6Pkt()
            Icmpv6Info = GetIcmpv6Info(protocolInfoList)
            Icmpv6.time = Icmpv6Info[0]
            Icmpv6.Icmpv6ND_RA = Icmpv6Info[1]
            Icmpv6.Icmpv6NDOptSrcLLAddr = Icmpv6Info[2]
            Icmpv6.Icmpv6NDOptMTU = Icmpv6Info[3]
            Icmpv6.Icmpv6NDOptPrefixInfo = Icmpv6Info[4]
            #Icmpv6.PrintMyInfo()
            pkt.icmpv6 = Icmpv6
        
        elif '<ICMP' in protocolInfoList[2]:
            pkt.flag['icmpv4'] = 1
            Icmpv4 = Icmpv4Pkt()
            Icmpv4Info = GetIcmpv4Info(protocolInfoList)
            Icmpv4.time = Icmpv4Info[0]
            Icmpv4.type_ = Icmpv4Info[1]
            Icmpv4.code = Icmpv4Info[2]
            Icmpv4.chksum = Icmpv4Info[3]
            Icmpv4.reserverd = Icmpv4Info[4]
            Icmpv4.length = Icmpv4Info[5]
            Icmpv4.nexthopmtu = Icmpv4Info[6]
            Icmpv4.unused = Icmpv4Info[7]
            Icmpv4.IpError = Icmpv4Info[8]
            Icmpv4.UdpError = Icmpv4Info[9]
            #Icmpv4.PrintMyInfo()
            pkt.icmpv4 = Icmpv4
        
    else:
        return